import torch
import random
import argparse
import contextlib
import numpy as np

from rich import print
from rich.text import Text
from rich.panel import Panel
from rich.prompt import Prompt

from gemma import config
from gemma import model as gemma_model

@contextlib.contextmanager
def _set_default_tensor_type(dtype: torch.dtype):
    """Sets the default torch dtype to the given dtype."""
    torch.set_default_dtype(dtype)
    yield
    torch.set_default_dtype(torch.float)


def main(args):
    model_config = config.get_model_config(args.variant)
    model_config.dtype = "float32" if args.device == "cpu" else "float16"
    model_config.quant = args.quant
    reply_size = args.reply_size
    
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    device = torch.device(args.device)
    with _set_default_tensor_type(model_config.get_dtype()):
        model = gemma_model.GemmaForCausalLM(model_config)
        model.load_weights(args.ckpt)
        model = model.to(device).eval()
    print("Model loading done")

    while True:
        prompt = Prompt.ask(">>> ")
        if prompt == "exit":
            break
        result = model.generate(prompt, device, reply_size)
        result = result.split("\n")[0]
        generated_text = Text(result, style="green")
        print(Panel(generated_text, width=32, title="Reply"))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", type=str, required=True)
    parser.add_argument("--variant",
                        type=str,
                        default="2b",
                        choices=["2b", "7b"])
    parser.add_argument("--device",
                        type=str,
                        default="cpu",
                        choices=["cpu", "cuda"])
    parser.add_argument("--seed", type=int, default=12345)
    parser.add_argument("--quant", action='store_true')
    parser.add_argument("--prompt", type=str, default="The meaning of life is")
    parser.add_argument("--reply-size", type=int, default=32)
    args = parser.parse_args()
    main(args)
